from typing import List, Optional
from fastapi import APIRouter, Depends, HTTPException, Query
from sqlalchemy.orm import Session
from sqlalchemy import and_, or_
from app.api.deps import get_db, get_current_user
from app.models.user import User
from app.models.ai_workflow import AIWorkflow, TaskStatus
from app.crud.ai_workflow import ai_workflow
from pydantic import BaseModel
from datetime import datetime

router = APIRouter()

class ImageTaskInfo(BaseModel):
    """图片任务信息"""
    workflow_id: int
    task_name: str
    status: str
    progress: int
    current_step: Optional[str] = None
    error_message: Optional[str] = None
    variation_count: int
    prompt: Optional[str] = None
    created_at: datetime
    updated_at: datetime
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None

class ImageTaskStatusResponse(BaseModel):
    """图片任务状态响应"""
    image_url: str
    tasks: List[ImageTaskInfo]
    total_tasks: int
    active_tasks: int  # 正在执行的任务数
    completed_tasks: int  # 已完成的任务数
    failed_tasks: int  # 失败的任务数

@router.get("/image-tasks/by-url", response_model=ImageTaskStatusResponse)
def get_image_tasks_by_url(
    *,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
    image_url: str = Query(..., description="图片URL")
):
    """根据图片URL获取相关的所有任务状态"""
    
    # 清理图片URL（去除可能的引号和空格）
    clean_image_url = image_url.strip().strip('"').strip("'").strip('`')
    
    # 查询包含该图片URL的所有工作流
    # 使用更宽松的匹配模式来处理可能的空格和格式差异
    workflows = db.query(AIWorkflow).filter(
        and_(
            AIWorkflow.user_id == current_user.id,
            or_(
                AIWorkflow.input_data.contains({"image_url": clean_image_url}),
                AIWorkflow.input_data.contains({"image_url": f" {clean_image_url} "}),
                AIWorkflow.input_data.like(f'%"image_url":%"{clean_image_url}"%'),
                AIWorkflow.input_data.like(f'%"image_url":%" {clean_image_url} "%')
            )
        )
    ).order_by(AIWorkflow.created_at.desc()).all()
    
    if not workflows:
        return ImageTaskStatusResponse(
            image_url=clean_image_url,
            tasks=[],
            total_tasks=0,
            active_tasks=0,
            completed_tasks=0,
            failed_tasks=0
        )
    
    # 构建任务信息列表
    tasks = []
    active_count = 0
    completed_count = 0
    failed_count = 0
    
    for workflow in workflows:
        input_data = workflow.input_data or {}
        
        task_info = ImageTaskInfo(
            workflow_id=workflow.id,
            task_name=workflow.task_name,
            status=workflow.status.value,
            progress=workflow.progress,
            current_step=workflow.current_step,
            error_message=workflow.error_message,
            variation_count=input_data.get("variation_count", 1),
            prompt=input_data.get("user_prompt") or input_data.get("prompt"),
            created_at=workflow.created_at,
            updated_at=workflow.updated_at,
            started_at=workflow.started_at,
            completed_at=workflow.completed_at
        )
        
        tasks.append(task_info)
        
        # 统计任务状态
        if workflow.status in [TaskStatus.PENDING, TaskStatus.RUNNING, TaskStatus.WAITING]:
            active_count += 1
        elif workflow.status == TaskStatus.COMPLETED:
            completed_count += 1
        elif workflow.status == TaskStatus.FAILED:
            failed_count += 1
    
    return ImageTaskStatusResponse(
        image_url=clean_image_url,
        tasks=tasks,
        total_tasks=len(tasks),
        active_tasks=active_count,
        completed_tasks=completed_count,
        failed_tasks=failed_count
    )

@router.get("/image-tasks/tasks/{task_id}")
def get_task_status(
    *,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
    task_id: int
):
    """获取特定任务的详细状态（兼容原有接口）"""
    
    workflow = ai_workflow.get_workflow(db=db, workflow_id=task_id)
    if not workflow:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    if workflow.user_id != current_user.id:
        raise HTTPException(status_code=403, detail="无权访问此任务")
    
    input_data = workflow.input_data or {}
    result_data = workflow.result_data or {}
    
    return {
        "id": workflow.id,
        "task_name": workflow.task_name,
        "status": workflow.status.value,
        "progress": workflow.progress,
        "current_step": workflow.current_step,
        "error_message": workflow.error_message,
        "input_data": input_data,
        "result_data": result_data,
        "created_at": workflow.created_at,
        "updated_at": workflow.updated_at,
        "started_at": workflow.started_at,
        "completed_at": workflow.completed_at,
        "original_image_url": input_data.get("image_url"),
        "variation_count": input_data.get("variation_count", 1),
        "prompt": input_data.get("user_prompt") or input_data.get("prompt")
    }